import numpy as np
from tqdm import tqdm
from copy import deepcopy
import xgboost as xgb
from utils import get_combination


class SelectFeature:
    """
      特征工程,特征数目:num\n
      小于等于num的特征数量则会采取所有特征的全组合以获取全局最优特征组合,
      否则，使用贪心算法获取特征组合
      """

    def __init__(self, model, score_fun, method='all'):
        self._defaults = ['greedy', 'all', 'tree']
        assert method in self._defaults, "method should in {}".format(self._defaults)

        self.model = model
        self.score_fun = score_fun

        self.features_name = None
        self.features_num = -1

        self.x = None
        self.y = None

        self.dtypes = None

        self.method = method

    def _set_params(self, x, y, feature):
        self.x = x
        self.y = y
        self.features_name = feature

        self.features_num = x.shape[0]

        self.dtypes = str(y.dtype)

    def _target(self, features):
        clf = self.model.fit(self.x[:, features], self.y)
        pre = clf.predict(self.x[:, features])
        score = self.score_fun(self.y, pre)
        return score

    def _all_select(self):
        lst = list(range(self.features_num))
        layers = get_combination(lst)

        best_score = -1
        best_features = None
        for i, layer in enumerate(layers):
            lst_layer = list(layer)
            print("第{}层:".format(i))
            for it in tqdm(lst_layer):
                score = self._target(list(it))
                if score > best_score:
                    best_score = score
                    best_features = it

        return best_features

    def _greedy_select(self):
        features = list(range(self.features_num))
        selected = []
        best_score = -1
        for i in range(self.features_num):
            best_feature = -1
            for it in features:
                buf = deepcopy(selected)
                buf.append(it)
                score = self._target(buf)
                if score > best_score:
                    best_score = score
                    best_feature = it

            if best_feature != -1:
                selected.append(best_feature)
        return selected

    def _tree_select(self):
        if self.dtypes[:3] == 'int':
            tree = xgb.XGBClassifier()
        else:
            tree = xgb.XGBRegressor()

        m = tree.fit(self.x, self.y)

        importance = m.feature_importance
        return importance

    def fit(self, x, y=None, features=None):
        self._set_params(x, y, features)
        fun = getattr(self, "_"+self.method+"_select")
        lst = fun()

        return lst
